"""
手动实现K近邻算法（K-Nearest Neighbors, KNN）
只使用Python标准库
"""

import math


def knn_predict(X_train, y_train, x_test, k=3, weighted=True):
    """
    K近邻预测（分类）- 简洁版本
    
    参数:
        X_train: 训练特征 [[x1, x2, ...], ...]
        y_train: 训练标签 [label1, label2, ...]
        x_test: 测试点 [x1, x2, ...] 或测试点列表 [[x1, x2, ...], ...]
        k: 最近邻数量
        weighted: 是否使用距离加权投票（True=加权，False=简单投票）
    
    返回:
        预测结果（单个值或列表）
    """
    # 处理单个测试点
    if isinstance(x_test[0], (int, float)):
        return _knn_single(X_train, y_train, x_test, k, weighted)
    
    # 批量预测
    return [_knn_single(X_train, y_train, x, k, weighted) for x in x_test]


def _knn_single(X_train, y_train, x, k, weighted):
    """单个测试点的KNN预测"""
    # 计算距离并排序，取前k个
    neighbors = sorted(
        [(math.sqrt(sum((a - b) ** 2 for a, b in zip(x, xi))), yi) 
         for xi, yi in zip(X_train, y_train)],
        key=lambda d: d[0]
    )[:k]
    
    if weighted:
        # 距离加权投票：权重 = 1/(距离+1e-10) 避免除零
        votes = {}
        for dist, label in neighbors:
            weight = 1.0 / (dist + 1e-10)
            votes[label] = votes.get(label, 0) + weight
        return max(votes.items(), key=lambda x: x[1])[0]
    else:
        # 简单投票：选择k个最近邻中最常见的标签
        labels = [label for _, label in neighbors]
        return max(set(labels), key=labels.count)


def knn_regress(X_train, y_train, x_test, k=3, weighted=True):
    """
    
    
    返回:
        预测值（单个值或列表）
    """
    # 处理单个测试点
    if isinstance(x_test[0], (int, float)):
        return _knn_regress_single(X_train, y_train, x_test, k, weighted)
    
    # 批量预测
    return [_knn_regress_single(X_train, y_train, x, k, weighted) for x in x_test]


def _knn_regress_single(X_train, y_train, x, k, weighted):
    """单个测试点的KNN回归"""
    # 计算距离并排序，取前k个
    neighbors = sorted(
        [(math.sqrt(sum((a - b) ** 2 for a, b in zip(x, xi))), yi) 
         for xi, yi in zip(X_train, y_train)],
        key=lambda d: d[0]
    )[:k]
    
    if weighted:
        # 距离加权平均：权重 = 1/(距离+1e-10)
        weights = [1.0 / (dist + 1e-10) for dist, _ in neighbors]
        values = [val for _, val in neighbors]
        return sum(w * v for w, v in zip(weights, values)) / sum(weights)
    else:
        # 简单平均
        return sum(val for _, val in neighbors) / k


# 测试代码
if __name__ == "__main__":
    print("=" * 60)
    print("K近邻算法（简洁版）")
    print("=" * 60)
    
    # 分类示例
    print("\n【分类示例】")
    X_train = [[1, 1], [1.5, 1.5], [2, 2], [5, 5], [5.5, 5.5], [6, 6]]
    y_train = [0, 0, 0, 1, 1, 1]
    X_test = [[1.2, 1.2], [5.3, 5.3], [3.5, 3.5]]
    
    print("训练数据:", list(zip(X_train, y_train)))
    print("测试数据:", X_test)
    
    # 加权投票
    preds_weighted = knn_predict(X_train, y_train, X_test, k=3, weighted=True)
    print("加权投票预测:", preds_weighted)
    
    # 简单投票
    preds_simple = knn_predict(X_train, y_train, X_test, k=3, weighted=False)
    print("简单投票预测:", preds_simple)
    
    # 单个点预测
    single_pred = knn_predict(X_train, y_train, [2, 2], k=3)
    print("单点预测 [2,2]:", single_pred)
    
    # 回归示例
    print("\n【回归示例】")
    X_train_reg = [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]]
    y_train_reg = [2.1, 4.0, 5.9, 8.2, 10.1]
    X_test_reg = [[1.5, 1.5], [3.5, 3.5]]
    
    print("训练数据:", list(zip(X_train_reg, y_train_reg)))
    print("测试数据:", X_test_reg)
    
    # 加权平均
    reg_weighted = knn_regress(X_train_reg, y_train_reg, X_test_reg, k=3, weighted=True)
    print("加权平均预测:", [f"{x:.2f}" for x in reg_weighted])
    
    # 简单平均
    reg_simple = knn_regress(X_train_reg, y_train_reg, X_test_reg, k=3, weighted=False)
    print("简单平均预测:", [f"{x:.2f}" for x in reg_simple])
    
    print("\n" + "=" * 60)